# -*- coding:utf8 -*-
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.linear_model import HuberRegressor
from sklearn.linear_model import Ridge
from sklearn.linear_model import Lasso
from sklearn.linear_model import ElasticNet
import numpy as np

# 线性回归练习
X_TRAIN = np.array([1, 2, 3, 3.5, 4, 5, 6, 9, 8.5])
Y_TRAIN = np.array([78, 87, 69, 75, 88, 92, 89, 98, 80])
X_TEST = 7.8


class Regression:
    def __init__(self, name, x_train, y_train):
        self.linear = None
        self.name = name
        self.x_train = x_train
        self.y_train = y_train

    def fit(self):
        if self.linear is not None:
            self.linear.fit(self.x_train.reshape(-1, 1), self.y_train)

    def predict(self, x):
        if self.linear is not None:
            return int(self.linear.predict(np.array([x]).reshape(-1, 1))[0])
        else:
            return None

    def predict_and_show(self, x):
        result = self.predict(x)
        print("对于测试数据{}，{}模型预测的结果是{}".format(x, self.name, result))

    def score(self):
        return self.linear.score(self.x_train.reshape(-1, 1), self.y_train)

    def score_and_show(self):
        score = self.score()
        print("{}模型的评分(R2 Score)为:{}".format(self.name, score))

    def draw(self):
        """
        得到最大最小训练值对应的响应，根据两个点画出拟合直线
        :return:
        """
        x_array = [self.x_train.min(), self.x_train.max()]
        y_array = list(map(lambda x: self.predict(x), x_array))
        plt.plot(x_array, y_array)


class MyL2LinearRegression(Regression):
    def __init__(self, x_train, y_train):
        super().__init__("L2线性回归", x_train, y_train)
        self.linear = LinearRegression()
        self.fit()


class MyHuberRegression(Regression):
    def __init__(self, x_train, y_train):
        super().__init__("Huber线性回归", x_train, y_train)
        self.linear = HuberRegressor()
        self.fit()


class MyRidgeRegression(Regression):
    def __init__(self, x_train, y_train):
        super().__init__("Ridge岭回归", x_train, y_train)
        self.linear = Ridge(10)
        self.fit()


class MyLassoRegression(Regression):
    def __init__(self, x_train, y_train):
        super().__init__("Lasso回归", x_train, y_train)
        self.linear = Lasso(10)
        self.fit()


class MyElasticNetRegression(Regression):
    def __init__(self, x_train, y_train):
        super().__init__("ElasticNet弹性网络回归", x_train, y_train)
        self.linear = ElasticNet(10, 0.3)
        self.fit()


# 画出散点图
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.scatter(X_TRAIN, Y_TRAIN, label="Training Sample"), plt.xlabel("学习时长(小时)"), plt.ylabel("学生成绩")


def work(regression):
    regression.predict_and_show(X_TEST)
    regression.draw()
    regression.score_and_show()


# 进行L2线性回归
work(MyL2LinearRegression(X_TRAIN, Y_TRAIN))
# 进行Huber线性回归
work(MyHuberRegression(X_TRAIN, Y_TRAIN))
# 进行岭回归
work(MyRidgeRegression(X_TRAIN, Y_TRAIN))
# 进行Lasso回归
work(MyLassoRegression(X_TRAIN, Y_TRAIN))
# 进行ElasticNet弹性网络回归
work(MyElasticNetRegression(X_TRAIN, Y_TRAIN))

plt.show()
